{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from pathlib import Path\n",
    "SCRATCH = Path.home() / \"scratch\"\n",
    "os.environ[\"HF_HOME\"] = str(SCRATCH / \"hf_home\")\n",
    "\n",
    "from transformers import AutoTokenizer\n",
    "\n",
    "CHECKPOINT_OR_NAME = 'McGill-NLP/nano-aha-moment-3b'\n",
    "# CHECKPOINT_OR_NAME = 'Qwen/Qwen2.5-3B'\n",
    "# CHECKPOINT_OR_NAME = 'google/gemma-3-4b-pt'\n",
    "# CHECKPOINT_OR_NAME = \"meta-llama/Llama-3.2-3B\"\n",
    "CHAT_MODEL_NAME = \"Qwen/Qwen2.5-3B-Instruct\" # should have the tokenizer we trained the checkpoint with\n",
    "# CHAT_MODEL_NAME = \"google/gemma-3-4b-it\"\n",
    "# CHAT_MODEL_NAME = \"meta-llama/Llama-3.2-3B-Instruct\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(CHAT_MODEL_NAME)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from vllm import LLM, SamplingParams\n",
    "\n",
    "inference_engine = LLM(\n",
    "    model=CHECKPOINT_OR_NAME,\n",
    "    gpu_memory_utilization=0.5,\n",
    "    dtype=torch.bfloat16, \n",
    "    swap_space=2,\n",
    "    enable_prefix_caching=True,\n",
    "    max_model_len=2048,\n",
    "    max_seq_len_to_capture=2048,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_response(query, response):\n",
    "    from IPython.display import HTML\n",
    "\n",
    "    # Escape <think> </think> <answer> </answer> and any HTML tags\n",
    "    response = response.replace(\"<\", \"&lt;\").replace(\">\", \"&gt;\")\n",
    "    query = query.replace(\"<\", \"&lt;\").replace(\">\", \"&gt;\")\n",
    "\n",
    "    # Format the response with syntax highlighting\n",
    "    formatted_html = f\"\"\"\n",
    "    <div style=\"background-color: #f8f9fa; padding: 15px; border-radius: 5px; border: 1px solid #ddd;\">\n",
    "        <h3 style=\"color: #333; margin-top: 0;\">Query:</h3>\n",
    "        <pre style=\"background-color: #e9f7fe; padding: 10px; border-radius: 3px; overflow-x: auto; white-space: pre-wrap; word-wrap: break-word; color: #0066cc;\">{query}</pre>\n",
    "        <h3 style=\"color: #333; margin-top: 10px;\">Generated Response:</h3>\n",
    "        <pre style=\"background-color: #f5f5f5; padding: 10px; border-radius: 3px; overflow-x: auto; white-space: pre-wrap; word-wrap: break-word; color: #333333;\">{response}</pre>\n",
    "    </div>\n",
    "    \"\"\"\n",
    "\n",
    "    return HTML(formatted_html)\n",
    "\n",
    "\n",
    "\n",
    "def generate_chat_prompt(query, assistance_prefix=\"Let me think step by step\\n<think>\"):\n",
    "    SYSTEM_MESSAGE = (\n",
    "        \"You are a helpful assistant. You first think about the reasoning process in the mind \"\n",
    "        \"and then provides the user with the answer.\"\n",
    "    )\n",
    "    r1_prefix = [\n",
    "        {\n",
    "            \"role\": \"system\",\n",
    "            \"content\": SYSTEM_MESSAGE,\n",
    "        },\n",
    "        {\"role\": \"user\", \"content\": f\"{query}\"},\n",
    "    ]\n",
    "    if assistance_prefix is not None:\n",
    "        r1_prefix.append({\"role\": \"assistant\", \"content\": assistance_prefix})\n",
    "    \n",
    "    input_ids = tokenizer.apply_chat_template(\n",
    "        r1_prefix, tokenize=True, continue_final_message=assistance_prefix is not None\n",
    "    )\n",
    "    prompt = tokenizer.decode(\n",
    "        input_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False\n",
    "    )\n",
    "    return {\"prompt\": prompt, \"input_ids\": input_ids}\n",
    "\n",
    "# play the countdown game\n",
    "def preprocess_countdown_example(example):\n",
    "    SYSTEM_MESSAGE = (\n",
    "        \"You are a helpful assistant. You first think about the reasoning process in the mind \"\n",
    "        \"and then provides the user with the answer.\"\n",
    "    )\n",
    "\n",
    "    PROMPT_TEMPLATE = (\n",
    "        \"Using the numbers {numbers}, create an equation that equals {target}. \"\n",
    "        \"You can use basic arithmetic operations (+, -, *, /) and each number can only be used once. \"\n",
    "        \"Show your work in <think> </think> tags. And return the final equation and answer in \"\n",
    "        \"<answer> </answer> tags, for example <answer>(1 + 2) / (3 * 5)</answer>.\"\n",
    "    )\n",
    "    numbers = example[\"nums\"]\n",
    "    target = example[\"target\"]\n",
    "    \n",
    "    chat_messages = [\n",
    "        {\"role\": \"system\",  \"content\": SYSTEM_MESSAGE},\n",
    "        {\"role\": \"user\", \"content\": PROMPT_TEMPLATE.format(numbers=numbers, target=target)},\n",
    "        {\"role\": \"assistant\", \"content\": \"Let me think step by step\\n<think>\"}\n",
    "    ]\n",
    "    \n",
    "    input_ids = tokenizer.apply_chat_template(\n",
    "        chat_messages, tokenize=True, continue_final_message=True\n",
    "    )\n",
    "    prompt = tokenizer.decode(input_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False)\n",
    "    \n",
    "    return {\"input_ids\": input_ids, \"prompt\": prompt}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Countdown"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample = {\"nums\":  [7, 71, 19, 4], \"target\": 68}\n",
    "sample.update(preprocess_countdown_example(sample))\n",
    "\n",
    "print(f\"######################## Prompt:\\n`{sample['prompt']}`\")\n",
    "\n",
    "generation = inference_engine.generate(\n",
    "    prompt_token_ids=sample[\"input_ids\"], \n",
    "    sampling_params=SamplingParams(\n",
    "        temperature=0.3,\n",
    "        max_tokens=1024,\n",
    "        top_p=1.0,\n",
    "        n=1,  # Only generate one response per question\n",
    "    )\n",
    ")\n",
    "response = tokenizer.decode(generation[0].outputs[0].token_ids)\n",
    "format_response(sample[\"prompt\"], response)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## General"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "o_1 = \"\"\"A quadratic equation has roots that are also the solutions to the system:\n",
    "\n",
    "\\[\n",
    "\\begin{cases}\n",
    "x + y = 7 \\\\\n",
    "xy = 10\n",
    "\\end{cases}\n",
    "\\]\n",
    "\n",
    "1. Find the quadratic equation whose roots are \\( x \\) and \\( y \\).\n",
    "2. Solve the quadratic equation.\n",
    "3. Verify that the roots satisfy the original system.\n",
    "\"\"\"\n",
    "\n",
    "o = \"\"\"Hello, how are you?\"\"\"\n",
    "\n",
    "o_3 = \"\"\"My slurm job failed. I look at the stdout and I observe this:\n",
    "\n",
    "### stdout ###\n",
    "slurmstepd: error: container_p_join: open failed for /var/opt/slurm/localstorage/6450599/.ns: No such file or directory\n",
    "slurmstepd: error: container_g_join(6450599): No such file or directory\"\"\"\n",
    "\n",
    "o_4 = \"\"\"\n",
    "How many of the first 500 positive integers are divisible by 3, 4 and 5?\n",
    "\"\"\"\n",
    "\n",
    "o_5 = \"\"\"\n",
    "You have 1,2,3,4. Provide a math equation that equals 10. You can use each number only once. You can use basic arithmetic operations (+, -, *, /).\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "sample = generate_chat_prompt((\n",
    "    f\"{o}\\n\"\n",
    "    \"Show your work in <think> </think> tags. And return the final answer in \"\n",
    "    \"<answer> </answer> tags\"\n",
    "))\n",
    "\n",
    "print(sample)\n",
    "\n",
    "print(f\"######################## Prompt:\\n`{sample['prompt']}`\")\n",
    "\n",
    "generation = inference_engine.generate(\n",
    "    prompt_token_ids=sample[\"input_ids\"], \n",
    "    sampling_params=SamplingParams(\n",
    "        temperature=0.6,\n",
    "        max_tokens=1024,\n",
    "        top_p=1.0,\n",
    "        n=1,  # Only generate one response per question\n",
    "    )\n",
    ")\n",
    "response = tokenizer.decode(generation[0].outputs[0].token_ids)\n",
    "format_response(o, response)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
